# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import collections
import logging

from absl import flags
import pandas
from coin.proto.coin_order_gateway_pb2 import OrderEvent
from coin.proto.coin_order_enums_pb2 import BUY_ORDER, SELL_ORDER
from coin.strategy.mm.tool.util import get_order_sign
import coin.proto.coin_market_enums_pb2 as cenums


class OrderQtyTrackerPerSide(object):
  def __init__(self):
    self._order_qty_map = {}

  def get_diff_qty(self, order_id, event):
    diff_qty = None
    if event.type == OrderEvent.ORDER_FILLED:
      diff_qty = -event.fill_qty
    elif event.type == OrderEvent.ORDER_ACCEPTED:
      diff_qty = event.order_qty
    elif event.type in [OrderEvent.CANCEL_CONFIRMED]:
      if order_id in self._order_qty_map:
        diff_qty = -self._order_qty_map[order_id][1]
    return diff_qty

  def get_diff_qty_from_json(self, order_id, event_json):
    diff_qty = None
    type = event_json['type']
    if type == 'ORDER_FILLED':
      diff_qty = -event_json['fill_qty']
    elif type == 'ORDER_ACCEPTED':
      diff_qty = event_json['order_qty']
    elif type in ['CANCEL_CONFIRMED']:
      if order_id in self._order_qty_map:
        diff_qty = -self._order_qty_map[order_id][1]
    return diff_qty

  def has_event(self, order_id):
    return order_id in self._order_qty_map

  def add_event(self, order_id, event):
    diff_qty = self.get_diff_qty(order_id, event)
    if diff_qty is None:
      return
    if order_id not in self._order_qty_map:
      self._order_qty_map[order_id] = [event.event_time, diff_qty]
    else:
      self._order_qty_map[order_id][1] += diff_qty
    if self._order_qty_map[order_id][1] < 1e-8:
      del self._order_qty_map[order_id]

  def add_event_json(self, order_id, event_json):
    diff_qty = self.get_diff_qty_from_json(order_id, event_json)
    if diff_qty is None:
      return
    if order_id not in self._order_qty_map:
      self._order_qty_map[order_id] = [event_json['event_time'], diff_qty]
    else:
      self._order_qty_map[order_id][1] += diff_qty
    if self._order_qty_map[order_id][1] < 1e-8:
      del self._order_qty_map[order_id]

  def get_order_existence(self):
    if len(self._order_qty_map) > 0:
      maxtime = max(self._order_qty_map.values())[0]
      removes = []
      for order_id, tq in self._order_qty_map.items():
        maxtime = int(maxtime)
        _time = int(tq[0])
        if maxtime - _time > 60e9:
          # logging.info("removing %s, elapsed %.2f sec", order_id, (maxtime - tq[0])*1e-9)
          removes.append(order_id)
      for remove in removes:
        del self._order_qty_map[remove]
    return len(self._order_qty_map) > 0


class OrderQtyTracker(object):
  def __init__(self):
    self._buy_track = OrderQtyTrackerPerSide()
    self._sell_track = OrderQtyTrackerPerSide()

  def add_event(self, order_id, event):
    sign = get_order_sign(event.order_side)
    if sign == 1:
      self._buy_track.add_event(order_id, event)
    elif sign == -1:
      self._sell_track.add_event(order_id, event)
    elif sign == 0:
      if self._buy_track.has_event(order_id):
        self._buy_track.add_event(order_id, event)
      elif self._sell_track.has_event(order_id):
        self._sell_track.add_event(order_id, event)
    else:
      raise ValueError(event)

  def add_event_json(self, order_id, event_json):
    if event_json['order_side'] in ['BUY_ORDER', 'BUY_OPEN_ORDER', 'BUY_CLOSE_ORDER']:
      self._buy_track.add_event_json(order_id, event_json)
    elif event_json['order_side'] in ['SELL_ORDER', 'SELL_OPEN_ORDER', 'SELL_CLOSE_ORDER']:
      self._sell_track.add_event_json(order_id, event_json)
    elif event_json['order_side'] == 'UNKNOWN_ORDER_SIDE':
      if self._buy_track.has_event(order_id):
        self._buy_track.add_event_json(order_id, event_json)
      elif self._sell_track.has_event(order_id):
        self._sell_track.add_event_json(order_id, event_json)
    else:
      raise ValueError(event_json)

  def get_order_existence(self):
    return (self._buy_track.get_order_existence(), self._sell_track.get_order_existence())


class DurationTracker(object):
  def __init__(self):
    self._oqtrack = collections.defaultdict(OrderQtyTracker)
    self._prev_buy_exist_time = collections.defaultdict(lambda: None)
    self._prev_sell_exist_time = collections.defaultdict(lambda: None)
    self._prev_timestamp = collections.defaultdict(lambda: None)
    self._buy_exist_duration = collections.defaultdict(int)
    self._sell_exist_duration = collections.defaultdict(int)
    self._all_duration = collections.defaultdict(int)
    self._buy_durs = collections.defaultdict(list)
    self._sell_durs = collections.defaultdict(list)

  def on_log(self, timestamp, event):
    # currently, other exchange types are not guaranteed to work well
    if event.exchange_type not in [cenums.Okex]:
      return
    if event.external_order_id == '':
      return
    symbol = event.symbol
    self._oqtrack[symbol].add_event(event.external_order_id, event)
    self.on_log_impl(timestamp, symbol)

  def on_log_json(self, timestamp, event_json):
    # currently, other exchange types are not guaranteed to work well
    if 'exchange_type' not in event_json or \
       event_json['exchange_type'] not in ['Okex']:
      return
    if 'external_order_id' not in event_json:
      return
    symbol = event_json['symbol']
    self._oqtrack[symbol].add_event_json(event_json['external_order_id'], event_json)
    self.on_log_impl(timestamp, symbol)

  def on_log_impl(self, timestamp, symbol):
    buy_exist, sell_exist = self._oqtrack[symbol].get_order_existence()
    if self._prev_timestamp[symbol] is not None:
      self._all_duration[symbol] += timestamp - self._prev_timestamp[symbol]
    self._prev_timestamp[symbol] = timestamp
    if buy_exist:
      if self._prev_buy_exist_time[symbol] is not None:
        self._buy_exist_duration[symbol] += timestamp - self._prev_buy_exist_time[symbol]
        self._buy_durs[symbol].append((timestamp, timestamp - self._prev_buy_exist_time[symbol]))
      self._prev_buy_exist_time[symbol] = timestamp
    else:
      self._prev_buy_exist_time[symbol] = None
    if sell_exist:
      if self._prev_sell_exist_time[symbol] is not None:
        self._sell_exist_duration[symbol] += timestamp - self._prev_sell_exist_time[symbol]
        self._sell_durs[symbol].append((timestamp, timestamp - self._prev_sell_exist_time[symbol]))
      self._prev_sell_exist_time[symbol] = timestamp
    else:
      self._prev_sell_exist_time[symbol] = None

  def report(self, symbol, need_plot):
    if self._all_duration[symbol] == 0:
      return
    buydur = pandas.Timedelta(self._buy_exist_duration[symbol])
    selldur = pandas.Timedelta(self._sell_exist_duration[symbol])
    totdur = pandas.Timedelta(self._all_duration[symbol] + 1)
    logging.info("%9s B-%s(%.2f) S-%s(%.2f) T-%s",
                 symbol,
                 buydur,
                 buydur / totdur,
                 selldur,
                 selldur / totdur,
                 totdur)
    if need_plot and len(self._buy_durs[symbol]) > 0 and len(self._sell_durs[symbol]) > 0:
      import matplotlib.pyplot as plt

      def plot_pss(durs):
        tss, diffs = zip(*durs)
        ys = []
        xmins = []
        xmaxs = []
        for ts, diff in zip(tss, diffs):
          ys.append(ts)
          xmins.append(ts)
          xmaxs.append(ts + diff)
        ys = pandas.DatetimeIndex(ys)
        xmins = pandas.DatetimeIndex(xmins)
        xmaxs = pandas.DatetimeIndex(xmaxs)
        plt.hlines(ys, xmins, xmaxs, lw=0.5)
        plt.xticks(rotation=40)

      plt.subplot(211)
      plot_pss(self._buy_durs[symbol])
      plt.subplot(212)
      plot_pss(self._sell_durs[symbol])
      plt.savefig("oe_stat/dur_%s_%s.png" % (symbol, flags.FLAGS.trading_date))
      plt.close()
